import cv2
import numpy as np

if __name__ == '__main__':
    # 1. 图片输入
    image_np = cv2.imread('Stephen_Curry.jpg')
    image_np1 = cv2.imread('Football_player.jpg')
    # 2. 灰度化
    image_np_gray = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY)
    image_np1_gray = cv2.cvtColor(image_np1, cv2.COLOR_BGR2GRAY)

    # 3. Canny算法
    image_edges = cv2.Canny(
        image_np_gray,
        30,
        70
    )

    # 4. 霍夫变换
    # 创建一张黑图
    drawing = np.zeros(image_np.shape, dtype=np.uint8)
    # 霍夫圆变换
    circles = cv2.HoughCircles(
        image_edges,  # 边缘检测图像
        cv2.HOUGH_GRADIENT,  # 霍夫梯度法
        1,  # dp
        20,  # 最小圆心距离
        param2=26,  # 累加器投票阈值
        minRadius=27,
        maxRadius=40
    )
    print(circles)  # 圆数据

    # 5. 图片输出
    # # 在黑图上画线
    for circle in circles[0]:
        print(circle)
        x, y, r = circle
        # 画圆
        cv2.circle(
            drawing,
            (int(x), int(y)),
            int(r),
            (0, 0, 255),
            2,
            lineType=cv2.LINE_AA
        )
        # 圆心
        cv2.circle(
            drawing,
            (int(x), int(y)),
            1,
            (0, 0, 255),
            2,
            lineType=cv2.LINE_AA
        )

    cv2.imshow('image_np', image_np)
    cv2.imshow('drawing', drawing)
    cv2.imwrite('basketball.jpg', drawing)
    # 3. Canny算法
    image_edges = cv2.Canny(
        image_np1_gray,
        30,
        70
    )

    # 4. 霍夫变换
    # 创建一张黑图
    drawing = np.zeros(image_np1.shape, dtype=np.uint8)
    # 霍夫圆变换
    circles = cv2.HoughCircles(
        image_edges,  # 边缘检测图像
        cv2.HOUGH_GRADIENT,  # 霍夫梯度法
        1,  # dp
        20,  # 最小圆心距离
        param2=26,  # 累加器投票阈值
        minRadius=50,
        maxRadius=60
    )
    print(circles)  # 圆数据

    # 5. 图片输出
    # # 在黑图上画线
    for circle in circles[0]:
        print(circle)
        x, y, r = circle
        # 画圆
        cv2.circle(
            drawing,
            (int(x), int(y)),
            int(r),
            (0, 0, 255),
            2,
            lineType=cv2.LINE_AA
        )
        # 圆心
        cv2.circle(
            drawing,
            (int(x), int(y)),
            1,
            (0, 0, 255),
            2,
            lineType=cv2.LINE_AA
        )

    cv2.imshow('image_np1', image_np1)
    cv2.imshow('drawing', drawing)
    cv2.imwrite('football.jpg', drawing)
    cv2.waitKey(0)

